import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Function
import torch.nn.functional as func
import numpy as np
from PIL import Image
import cv2
import matplotlib.pyplot as plt
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.optim import lr_scheduler

from epe_checker import report_epe_violations, get_epe_checkpoints
import lithosim_cuda as litho
import ilt_utils as ilt

# 频域低通滤波函数
def frequency_filter(mask, cutoff_frequency=18):
    # 对掩模进行2D FFT
    mask_fft = torch.fft.fft2(mask)

    # 获取掩模的尺寸
    rows, cols = mask.shape
    center_row, center_col = rows // 2, cols // 2

    # 在频域中保留低频部分，剪除高频部分
    mask_fft[:center_row - cutoff_frequency, :] = 0
    mask_fft[center_row + cutoff_frequency:, :] = 0
    mask_fft[:, :center_col - cutoff_frequency] = 0
    mask_fft[:, center_col + cutoff_frequency:] = 0

    return mask_fft


# 低频保持损失
def low_frequency_loss(original_mask, optimized_mask, cutoff_frequency=18):
    original_fft = torch.fft.fft2(original_mask)
    optimized_fft = torch.fft.fft2(optimized_mask)

    # 计算低频部分的L2损失
    original_fft_filtered = frequency_filter(original_mask, cutoff_frequency)
    optimized_fft_filtered = frequency_filter(optimized_mask, cutoff_frequency)

    loss = func.mse_loss(optimized_fft_filtered, original_fft_filtered)
    return loss


# 二值化损失
def binarization_loss(mask,target):
    return func.mse_loss(mask, target)


class ilt_loss_function(Function):
    r"""
    The forward/backward implementation of conventional ILT functionality
    """
    @staticmethod
    def forward(ctx, mask, target, kernels_f, kernels_flip_f, kernels_ct_f, kernels_ct_flip_f, kernels_df, kernels_flip_df, kernels_ct_df, kernels_ct_flip_df, weight_focus, weight_defocus, dose_params):
        r"""
        同时计算Nom/Max/Min三个剂量条件下的前向传播
        Args:
            dose_params: tuple包含三个剂量值 (dose_nom, dose_max, dose_min)
        """
        dose_nom, dose_max, dose_min = dose_params

        # 前向传播分支1: Nominal
        aerialNom = ilt_loss_function._single_dose_forward(mask, kernels_f, weight_focus, dose_nom)
        # 前向传播分支2: Max
        aerialMax = ilt_loss_function._single_dose_forward(mask, kernels_f, weight_focus, dose_max)
        # 前向传播分支3: Min
        aerialMin = ilt_loss_function._single_dose_forward(mask, kernels_df, weight_defocus, dose_min)

        # 保存反向传播所需参数
        ctx.save_for_backward(
            mask,
            kernels_f, kernels_flip_f, kernels_ct_f, kernels_ct_flip_f,
            kernels_df, kernels_flip_df, kernels_ct_df, kernels_ct_flip_df,
            weight_focus,weight_defocus
        )
        ctx.dose_params = dose_params

        return aerialNom, aerialMax, aerialMin

    @staticmethod
    def _single_dose_forward(mask, kernels, weight, dose):
        """单个剂量条件的前向计算"""
        cmask = litho.tensor_real_to_complex(mask, dose=dose)
        cmask = litho.fft2(cmask)
        cmask = litho.frequency_multiplication(cmask, kernels)
        cmask = litho.ifft2(cmask)
        return litho.tensor_weight_sum(cmask, weight, square_root=False, normalized_weight=False)

    @staticmethod
    def backward(ctx, grad_nom, grad_max, grad_min):
        """
        处理三个分支的梯度反向传播
        grad_nom/grad_max/grad_min: 对应三个分支的梯度输入
        """
        # 获取保存的参数
        (mask, kernels_f, kernels_flip_f, kernels_ct_f, kernels_ct_flip_f,
         kernels_df, kernels_flip_df, kernels_ct_df, kernels_ct_flip_df,
         weight_focus,weight_defocus) = ctx.saved_tensors

        dose_nom, dose_max, dose_min = ctx.dose_params

        # 计算总梯度 = grad_nom分量 + grad_max分量 + grad_min分量
        total_grad = 0.0

        # 分支1: Nominal的梯度贡献
        grad_nom_term = ilt_loss_function._single_dose_backward(
            mask, grad_nom,
            kernels_f, kernels_flip_f, kernels_ct_f, kernels_ct_flip_f,
            weight_focus, dose_nom
        )
        total_grad += grad_nom_term

        # 分支2: Max的梯度贡献
        grad_max_term = ilt_loss_function._single_dose_backward(
            mask, grad_max,
            kernels_f, kernels_flip_f, kernels_ct_f, kernels_ct_flip_f,
            weight_focus, dose_max
        )
        total_grad += grad_max_term

        # 分支3: Min的梯度贡献
        grad_min_term = ilt_loss_function._single_dose_backward(
            mask, grad_min,
            kernels_df, kernels_flip_df, kernels_ct_df, kernels_ct_flip_df,
            weight_defocus, dose_min
        )
        total_grad += grad_min_term

        # total_grad = ilt.ideal_lpfilter(mask_backup, 193, 1.35, 1)

        # 返回梯度 (None对应不需要梯度的参数)
        return (total_grad.real,
                None, None, None, None, None, None, None, None, None, None, None, None)

    @staticmethod
    def _single_dose_backward(mask, grad_output, kernels, kernels_flip, kernels_ct, kernels_ct_flip, weight, dose):
        """单个分支的反向传播计算"""
        # 复现前向传播中的中间步骤
        mask_ct_conv = litho.convolve_kernel(mask, kernels_ct, weight, dose=dose, combo_kernel=False)
        grad_left_term = litho.convolve_kernel(mask_ct_conv, kernels_flip, weight, dose=dose, combo_kernel=False)

        mask_conv = litho.convolve_kernel(mask, kernels, weight, dose=dose, combo_kernel=False)
        grad_right_term = litho.convolve_kernel(mask_conv, kernels_ct_flip, weight, dose=dose, combo_kernel=False)

        return grad_output * (grad_left_term + grad_right_term)


if __name__ == '__main__':
    # 设置字体为支持中文的字体，例如 SimHei（黑体）
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置字体为 SimHei
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 原尺寸核
    kernel_focus = torch.load('kernel_neuralilt/kernel_focus_tensor.pt', map_location=device)
    kernel_focus_ct = torch.load('kernel_neuralilt/kernel_ct_focus_tensor.pt', map_location=device)
    kernel_defocus = torch.load('kernel_neuralilt/kernel_defocus_tensor.pt', map_location=device)
    kernel_defocus_ct = torch.load('kernel_neuralilt/kernel_ct_defocus_tensor.pt', map_location=device)

    kernel_focus_flip = torch.flip(kernel_focus, [1,2])
    kernel_focus_ct_flip = torch.flip(kernel_focus_ct, [1, 2])
    kernel_defocus_flip = torch.flip(kernel_defocus, [1,2])
    kernel_defocus_ct_flip = torch.flip(kernel_defocus_ct, [1,2])
    # 权重
    weight_focus = torch.load('kernel_neuralilt/weight_focus_tensor.pt', map_location=device)
    weight_defocus = torch.load('kernel_neuralilt/weight_defocus_tensor.pt', map_location=device)

    # 加载并对mask初始化
    mask_path = 'ICCAD2013/png/target1.png'
    mask_backup = litho.load_image(mask_path) # [1,2048,2048]
    mask_backup = mask_backup.to(device)
    if len(mask_backup.shape) == 3:
        mask_backup = mask_backup.unsqueeze(0)  # [1,1,2048,2048] 源代码第一维为batch

    target = mask_backup.clone().detach().requires_grad_(False)
    target8 = func.avg_pool2d(target, 8, stride=8)
    target4 = func.avg_pool2d(target, 4, stride=4)
    target2 = func.avg_pool2d(target, 2, stride=2)

    # 一阶段初始化
    # mask_backup = ilt.bit_mask_to_two_value_mask(mask_backup)
    # mask_backup = ilt.gauss_lpfilter(mask_backup, 193, 1.35, 1, pixel_size=1)
    mask_backup = ilt.ideal_lpfilter(mask_backup, 193, 1.35, 1) #0-1->负数
    # plt.imshow(mask_backup[0, 0, :, :].cpu().numpy())
    # plt.title("原图低通")
    # plt.colorbar()
    # plt.show()
    mask_backup = ilt.relex_mask(mask_backup, theta_m=4, threhold_m=0.5) #复数->0-1
    # plt.imshow(mask_backup[0, 0, :, :].cpu().numpy())
    # plt.title("原图松弛")
    # plt.colorbar()
    # plt.show()

    # 一阶段缩小8倍优化
    mask = mask_backup.clone().detach().requires_grad_(True)
    mask_tmp8 = mask_backup.clone().detach().requires_grad_(True)
    mask8 = func.avg_pool2d(mask_tmp8, 8, stride=8)
    mask8 = mask8.clone().detach().requires_grad_(True)

    grad_values = []
    optimizer = optim.Adam([mask8], lr=0.1)

    # 保存最优结果
    lossMin, l2Min, pvbMin = 1e12, 1e12, 1e12
    best_mask = None
    # 定义剂量参数元组
    DOSE_PARAMS = (1.0, 1.02, 0.98)  # (nominal, max, min)
    WeightEPE = 0.5
    WeightPVBL2 = 1.0
    WeightPVBand = 0.0
    for epoch in range(100):
        # 单次前向传播获取三个剂量条件下的结果
        aerialNom, aerialMax, aerialMin = ilt_loss_function.apply(
            mask8,target8,
            kernel_focus,kernel_focus_flip,kernel_focus_ct,kernel_focus_ct_flip,
            kernel_defocus,kernel_defocus_flip,kernel_defocus_ct,kernel_defocus_ct_flip,
            weight_focus,weight_defocus,DOSE_PARAMS
        )

        printedNom = ilt.sigmoid_ilt_z(aerialNom, theta_z=50, threhold_z=0.225)
        printedMax = ilt.sigmoid_ilt_z(aerialMax, theta_z=50, threhold_z=0.225)
        printedMin = ilt.sigmoid_ilt_z(aerialMin, theta_z=50, threhold_z=0.225)

        l2loss = func.mse_loss(printedNom, target8, reduction="sum")
        pvbl2 = func.mse_loss(printedMax, target8, reduction="sum") + func.mse_loss(printedMin, target8, reduction="sum")
        pvbloss = func.mse_loss(printedMax, printedMin, reduction="sum")
        pvband = torch.sum((printedMax >= 0.225) != (printedMin >= 0.225))

        loss = l2loss + WeightPVBL2 * pvbl2

        print(f'原尺寸 Epoch {epoch + 1}, Loss: {loss.item():.4f}, L2: {l2loss.item():.4f}, PVB: {pvband.item():.4f}')

        # 记录最佳结果
        if best_mask is None or loss.item() < lossMin:
            lossMin, l2Min, pvbMin = loss.item(), l2loss.item(), pvband.item()
            best_mask = mask8.detach().clone()
            # plt.imshow(best_mask.squeeze().cpu().numpy())
            # plt.colorbar()
            # plt.show()

        optimizer.zero_grad()
        # loss.backward(retain_graph=True)
        loss.backward()
        # loss.requires__grad__(True)
        optimizer.step()

    plt.imshow(best_mask.squeeze().cpu().numpy())
    plt.title('缩小8倍掩模优化灰度结果')
    plt.colorbar()
    plt.show()
    best_relex= ilt.relex_mask(mask_backup, theta_m=4, threhold_m=0.5)
    plt.imshow(best_relex.squeeze().cpu().numpy())
    plt.title('缩小8倍掩模优化松弛结果')
    plt.colorbar()
    plt.show()

    ## 一阶段缩小4倍
    mask_tmp4 = best_mask.clone().detach().requires_grad_(True)
    mask4 = func.interpolate(mask_tmp4, scale_factor=2, mode="nearest")
    mask4 = mask4.clone().detach().requires_grad_(True)

    # 梯度下降设置
    grad_values = []
    optimizer = optim.Adam([mask4], lr=0.05)

    # 保存最优结果
    lossMin, l2Min, pvbMin = 1e12, 1e12, 1e12
    best_mask = None
    # 定义剂量参数元组
    DOSE_PARAMS = (1.0, 1.02, 0.98)  # (nominal, max, min)
    WeightEPE = 0.5
    WeightPVBL2 = 1.0
    WeightPVBand = 0.0
    for epoch in range(100):
        # 单次前向传播获取三个剂量条件下的结果
        aerialNom, aerialMax, aerialMin = ilt_loss_function.apply(
            mask4, target4,
            kernel_focus, kernel_focus_flip, kernel_focus_ct, kernel_focus_ct_flip,
            kernel_defocus, kernel_defocus_flip, kernel_defocus_ct, kernel_defocus_ct_flip,
            weight_focus, weight_defocus, DOSE_PARAMS
        )

        printedNom = ilt.sigmoid_ilt_z(aerialNom, theta_z=50, threhold_z=0.225)
        printedMax = ilt.sigmoid_ilt_z(aerialMax, theta_z=50, threhold_z=0.225)
        printedMin = ilt.sigmoid_ilt_z(aerialMin, theta_z=50, threhold_z=0.225)

        l2loss = func.mse_loss(printedNom, target4, reduction="sum")
        pvbl2 = func.mse_loss(printedMax, target4, reduction="sum") + func.mse_loss(printedMin, target4,
                                                                                    reduction="sum")
        pvbloss = func.mse_loss(printedMax, printedMin, reduction="sum")
        pvband = torch.sum((printedMax >= 0.225) != (printedMin >= 0.225))

        loss = l2loss + WeightPVBL2 * pvbloss

        print(f'原尺寸 Epoch {epoch + 1}, Loss: {loss.item():.4f}, L2: {l2loss.item():.4f}, PVB: {pvband.item():.4f}')

        # 记录最佳结果
        if best_mask is None or loss.item() < lossMin:
            lossMin, l2Min, pvbMin = loss.item(), l2loss.item(), pvband.item()
            best_mask = mask4.detach().clone()
            # plt.imshow(best_mask.squeeze().cpu().numpy())
            # plt.colorbar()
            # plt.show()

        optimizer.zero_grad()
        # loss.backward(retain_graph=True)
        loss.backward()
        # loss.requires__grad__(True)
        optimizer.step()

    plt.imshow(best_mask.squeeze().cpu().numpy())
    plt.title('缩小4倍掩模优化灰度结果')
    plt.colorbar()
    plt.show()

    ## 一阶段缩小2倍
    mask_tmp2 = best_mask.clone().detach().requires_grad_(True)
    mask2 = func.interpolate(mask_tmp2, scale_factor=2, mode="nearest")
    mask2 = mask2.clone().detach().requires_grad_(True)

    # 梯度下降设置
    grad_values = []
    optimizer = optim.Adam([mask2], lr=0.05)

    # 保存最优结果
    lossMin, l2Min, pvbMin = 1e12, 1e12, 1e12
    best_mask = None
    # 定义剂量参数元组
    DOSE_PARAMS = (1.0, 1.02, 0.98)  # (nominal, max, min)
    WeightEPE = 0.5
    WeightPVBL2 = 1.0
    WeightPVBand = 0.0
    for epoch in range(50):
        # 单次前向传播获取三个剂量条件下的结果
        aerialNom, aerialMax, aerialMin = ilt_loss_function.apply(
            mask2, target2,
            kernel_focus, kernel_focus_flip, kernel_focus_ct, kernel_focus_ct_flip,
            kernel_defocus, kernel_defocus_flip, kernel_defocus_ct, kernel_defocus_ct_flip,
            weight_focus, weight_defocus, DOSE_PARAMS
        )

        printedNom = ilt.sigmoid_ilt_z(aerialNom, theta_z=50, threhold_z=0.225)
        printedMax = ilt.sigmoid_ilt_z(aerialMax, theta_z=50, threhold_z=0.225)
        printedMin = ilt.sigmoid_ilt_z(aerialMin, theta_z=50, threhold_z=0.225)

        l2loss = func.mse_loss(printedNom, target2, reduction="sum")
        pvbl2 = func.mse_loss(printedMax, target2, reduction="sum") + func.mse_loss(printedMin, target2,
                                                                                    reduction="sum")
        pvbloss = func.mse_loss(printedMax, printedMin, reduction="sum")
        pvband = torch.sum((printedMax >= 0.225) != (printedMin >= 0.225))

        loss = l2loss + WeightPVBL2 * pvbloss

        print(f'原尺寸 Epoch {epoch + 1}, Loss: {loss.item():.4f}, L2: {l2loss.item():.4f}, PVB: {pvband.item():.4f}')

        # 记录最佳结果
        if best_mask is None or loss.item() < lossMin:
            lossMin, l2Min, pvbMin = loss.item(), l2loss.item(), pvband.item()
            best_mask = mask2.detach().clone()
            # plt.imshow(best_mask.squeeze().cpu().numpy())
            # plt.colorbar()
            # plt.show()

        optimizer.zero_grad()
        # loss.backward(retain_graph=True)
        loss.backward()
        # loss.requires__grad__(True)
        optimizer.step()

    plt.imshow(best_mask.squeeze().cpu().numpy())
    plt.title('缩小2倍掩模优化灰度结果')
    plt.colorbar()
    plt.show()




    ## 二阶段：将灰度掩模恢复成二值掩模
    ## 要求：优化掩模低频基本不变，高频细节变化，并将其逐渐转变为二值掩模，尽量不要用成像模型
    # 需要优化的掩模
    bmask_tmp = best_mask.clone().detach().requires_grad_(True)
    bmask = func.interpolate(bmask_tmp, scale_factor=2, mode="nearest")
    bmask = bmask.clone().detach().requires_grad_(True)
    bmask_t = bmask.clone().detach().requires_grad_(False)
    target_lf = ilt.ideal_lpfilter(bmask_t, 193, 1.3, 1)

    # 优化过程
    optimizer = torch.optim.Adam([bmask], lr=0.01)
    lossMin = 1e12

    for iteration in range(300):
        optimizer.zero_grad()

        # 计算损失：低频保持损失 + 二值化损失
        bmask_lf = ilt.ideal_lpfilter(bmask, 193, 1.3, 1)
        loss_low_freq = func.mse_loss(bmask_lf, target_lf)
        loss_bin = binarization_loss(bmask, target)

        # 总损失
        loss = 100 * loss_low_freq + 100 * loss_bin

        # 记录最佳结果
        if best_mask is None or loss.item() < lossMin:
            best_mask = bmask.detach().clone()
            # plt.imshow(best_mask.squeeze().cpu().numpy())
            # plt.colorbar()
            # plt.show()

        print(f"Iteration {iteration}, Loss: {loss.item()},loss_low_freq={loss_low_freq}, loss_bin={loss_bin}")

        # 反向传播和优化
        loss.backward()
        optimizer.step()

    print(f"结束")


    # # Save the best mask image
    # if best_mask is not None:
    #     # 将中间变量变成原本mask
    #     sigmod_mask = ilt.relex_mask(best_mask,theta_m=4,threhold_m=0.5)
    #     sigmod_mask = best_mask
    #     # 二值化掩模
    #     binary_mask = ilt.binary_mask(sigmod_mask, 0.5)
    #
    #     sig_best_intensity_map, sig_best_binary_wafer = litho.lithosim(sigmod_mask, 0.225, kernel_focus,
    #                                                                    weight_focus, None,
    #                                                                    False, return_binary_wafer=True)
    #     bi_best_intensity_map, bi_best_binary_wafer = litho.lithosim(binary_mask, 0.225, kernel_focus,
    #                                                                    weight_focus, None,
    #                                                                  False,return_binary_wafer=True)
    #
    #     plt.imshow(sigmod_mask.squeeze().cpu().numpy())
    #     plt.title('原尺寸掩模优化灰度结果')
    #     plt.colorbar()
    #     plt.show()
    #     plt.imshow(binary_mask.squeeze().cpu().numpy())
    #     plt.title('原尺寸掩模优化二值化结果')
    #     plt.colorbar()
    #     plt.show()
    #
    #     plt.imshow(sig_best_intensity_map.squeeze().cpu().numpy())
    #     plt.title('灰度原尺寸空间像')
    #     plt.colorbar()
    #     plt.show()
    #     plt.imshow(sig_best_binary_wafer.squeeze().cpu().numpy())
    #     plt.colorbar()
    #     plt.title('灰度原尺寸光刻胶像')
    #     plt.show()
    #     plt.imshow(bi_best_intensity_map.squeeze().cpu().numpy())
    #     plt.title('二值原尺寸空间像')
    #     plt.colorbar()
    #     plt.show()
    #     plt.imshow(bi_best_binary_wafer.squeeze().cpu().numpy())
    #     plt.colorbar()
    #     plt.title('二值原尺寸光刻胶像')
    #     plt.show()
    #
    #     # sigmod_mask去掉前两个维度
    #     sigmod_mask = sigmod_mask.squeeze(0).squeeze(0)
    #     cv2.imwrite(f"MOSAIC_gray.png", (sigmod_mask * 255).detach().cpu().numpy())
    #     binary_mask = binary_mask.squeeze(0).squeeze(0)
    #     cv2.imwrite(f"MOSAIC_binary.png", (binary_mask * 255).detach().cpu().numpy())



# class ilt_loss(torch.autograd.Function):
#     @staticmethod
#     def forward(ctx, theta, target, kernels, kernels_ct, weight):
#         gamma = 4
#         mask = (torch.cos(theta) + 1) / 2
#         mask_lf = ut.ideal_lpfilter(mask, 193, 1.35)
#         # target_lf = ut.ideal_lpfilter(target, 193, 1.35)
#         inom_lf_result, _ = litho.lithosim(mask_lf, 0.225, kernels, weight,
#                                    None, False, return_binary_wafer=False)
#         znom_lf_result = ut.sigmoid_ilt_z(inom_lf_result, theta_z=50, threhold_z=0.225)
#         result = (znom_lf_result - target).pow(gamma).sum()  # 空间像与traget的差值的4次方和
#         ctx.save_for_backward(theta, target, kernels, kernels_ct, weight)
#         return result
#
#     @staticmethod
#     def backward(ctx, grad_output):
#         theta, target, kernels, kernels_ct, weight = ctx.saved_tensors
#         grad_input = ut.compute_gradient(theta, target, kernels, kernels_ct, weight,
#                                          dose=1.0, gamma=4.0, theta_z=50, theta_m=4, avgpool_size=None)
#         grad_input = grad_input * 2
#
#         return grad_output * grad_input, None, None, None, None

# # 2025-2-3版本,检查梯度计算问题
# if __name__ == '__main__':
#     # 基础配置
#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#
#     # 原尺寸核
#     kernel_focus = torch.load('kernel_neuralilt/kernel_focus_tensor.pt', map_location=device)
#     kernel_focus_ct = torch.load('kernel_neuralilt/kernel_ct_focus_tensor.pt', map_location=device)
#     # 权重
#     weight_focus = torch.load('kernel_neuralilt/weight_focus_tensor.pt', map_location=device)
#
#
#     mask_path = 'ICCAD2013/png/target1.png'
#     mask = ut.load_image(mask_path)  # [1,1,2048,2048]
#     target = ut.load_image(mask_path)
#
#     # 设置字体为支持中文的字体，例如 SimHei（黑体）
#     plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置字体为 SimHei
#     plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
#
#     # Ensure mask has the correct dimensions
#     if len(mask.shape) == 3:
#         mask = mask.unsqueeze(0)  # Add batch dimension
#
#     grad_values = []
#     initial_weight = 100
#     decay_rate = 0.99
#
#     target_tmp = target.clone().detach().requires_grad_(False)
#     mask_tmp = mask.clone().detach().requires_grad_(False)
#     # 对mask进行初始化
#     mask_tmp = ut.sigmoid_ilt_mask(mask_tmp)
#
#     theta_back = torch.acos(2 * mask_tmp - 1)
#     theta = theta_back.clone().detach().requires_grad_(True)
#
#     optimizer = optim.Adam([theta], lr=0.01)
#     scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5, verbose=True)
#
#     best_loss = float('inf')
#     best_theta = None
#     best_intensity_map = None
#     best_binary_wafer = None
#
#     # 优化循环
#     print(f'原尺寸 分辨率为 {mask_tmp.shape[2]} x {mask_tmp.shape[3]}: ')
#     for epoch in range(300):
#         l2_loss = ilt_loss.apply(theta, target, kernel_focus, kernel_focus_ct, weight_focus)
#
#         theta_mask = (torch.cos(theta) + 1) / 2
#         aerial_mask = litho.lithosim(theta_mask, 0.225, kernel_focus,
#                                      weight_focus, None, False)
#         # 空间像惩罚项
#         # aerial_regu = torch.nn.functional.mse_loss(aerial_mask, target, reduction='sum')
#         aerial_weight = 0.1
#         # regu1 = torch.sum(theta_mask ** 2 * (theta_mask - 1) ** 2)
#         # regu1_weight = initial_weight * (decay_rate ** epoch)  # 指数衰减
#         # loss = l2_loss + regu1_weight * regu1
#
#         loss = l2_loss
#
#         print(f'原尺寸 Epoch {epoch + 1}, Loss: {loss.item():.4f}')
#
#         # 记录最佳结果
#         if best_loss is None or best_theta is None or loss.item() < best_loss:
#             best_loss = loss.detach().clone()
#             best_theta = theta.detach().clone()
#             # plt.imshow(best_mask3.squeeze().cpu().numpy())
#             # plt.show()
#
#         optimizer.zero_grad()
#         loss.backward()
#         # 记录梯度均值
#         grad_values.append(theta.grad.abs().mean().item())
#         optimizer.step()
#
#     # 绘制梯度变化趋势
#     plt.plot(grad_values)
#     plt.xlabel('Epoch')
#     plt.ylabel('Gradient Mean')
#     plt.title('Gradient Vanishing/Exploding Check')
#     plt.show()
#
#     # Save the best mask image
#     if best_theta is not None:
#         best_mask = (torch.cos(best_theta) + 1) / 2
#         # best_mask = litho.mask_sigmod_threshold(best_mask, steepness=50, sigmod_threshold=0.225)
#         best_mask = litho.mask_step_threshold(best_mask, 0.5)
#         best_intensity_map, _ = litho.lithosim(best_mask, 0.225, kernel_focus, weight_focus, None,
#                                                                False)
#         best_binary_wafer = ut.sigmoid_ilt_z(best_intensity_map, theta_z=50, threhold_z=0.225)
#         plt.imshow(best_mask.squeeze().cpu().numpy())
#         plt.title('原尺寸掩模优化结果')
#         plt.colorbar()
#         plt.show()
#
#     # Save the best intensity_map and binary_wafer
#     if best_intensity_map is not None and best_binary_wafer is not None:
#         plt.imshow(best_intensity_map.squeeze().cpu().numpy())
#         plt.title('原尺寸空间像')
#         plt.colorbar()
#         plt.show()
#         plt.imshow(best_binary_wafer.squeeze().cpu().numpy())
#         plt.colorbar()
#         plt.title('原尺寸光刻胶像')
#         plt.show()


# # 2025-1-28版本，梯度滤波和kerenl降采样
# if __name__ == '__main__':
#     # 基础配置
#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#
#     # 权重
#     weight_focus = torch.load('kernel_neuralilt/weight_focus_tensor.pt', map_location=device)
#     weight_defocus = torch.load('kernel_neuralilt/weight_defocus_tensor.pt', map_location=device)
#
#     # 原尺寸核
#     kernel_focus = torch.load('kernel_neuralilt/kernel_focus_tensor.pt', map_location=device)
#     kernel_focus_ct = torch.load('kernel_neuralilt/kernel_ct_focus_tensor.pt', map_location=device)
#     kernel_defocus = torch.load('kernel_neuralilt/kernel_defocus_tensor.pt', map_location=device)
#     kernel_defocus_ct = torch.load('kernel_neuralilt/kernel_ct_defocus_tensor.pt', map_location=device)
#
#     # 降采样4倍核
#     kernel_focus_lowres4 = torch.load('kernel_neuralilt/kernel_focus_tensor_lowres4.pt', map_location=device)
#     kernel_focus_ct_lowres4 = torch.load('kernel_neuralilt/kernel_ct_focus_tensor_lowres4.pt', map_location=device)
#     kernel_defocus_lowres4 = torch.load('kernel_neuralilt/kernel_defocus_tensor_lowres4.pt', map_location=device)
#     kernel_defocus_ct_lowres4 = torch.load('kernel_neuralilt/kernel_ct_defocus_tensor_lowres4.pt', map_location=device)
#
#     # 降采样2倍核
#     kernel_focus_lowres2 = torch.load('kernel_neuralilt/kernel_focus_tensor_lowres2.pt', map_location=device)
#     kernel_focus_ct_lowres2 = torch.load('kernel_neuralilt/kernel_ct_focus_tensor_lowres2.pt', map_location=device)
#     kernel_defocus_lowres2 = torch.load('kernel_neuralilt/kernel_defocus_tensor_lowres2.pt', map_location=device)
#     kernel_defocus_ct_lowres2 = torch.load('kernel_neuralilt/kernel_ct_defocus_tensor_lowres2.pt', map_location=device)
#
#     mask_path = 'ICCAD2013/png/target1.png'
#     mask_highres = ut.load_image(mask_path)  # [1,1,2048,2048]
#     target_highres = ut.load_image(mask_path)
#
#     # 降采样
#     avgpool_size = 4  # 降采样因子
#     avg_layer = nn.AvgPool2d(avgpool_size, stride=avgpool_size)
#     mask_lowres = avg_layer(mask_highres)  # [1,512,512]
#     target_lowres = avg_layer(target_highres)  # [1,512,512]
#
#     avgpool_size2 = 2  # 降采样因子
#     avg_layer2 = nn.AvgPool2d(avgpool_size2, stride=avgpool_size2)
#     target_lowres2 = avg_layer2(target_highres)
#
#     # 设置字体为支持中文的字体，例如 SimHei（黑体）
#     plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置字体为 SimHei
#     plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
#
#
#     # 1.先降采样，然后做梯度下降
#     mask1 = mask_lowres
#     target1 = target_lowres
#     target1 = target1.clone().detach().requires_grad_(False)
#
#     # Ensure mask has the correct dimensions
#     if len(mask1.shape) == 3:
#         mask1 = mask1.unsqueeze(0)  # Add batch dimension
#
#     # 使用cos松弛
#     mask1 = mask1.clone().detach().requires_grad_(False)
#     mask1 = ut.sigmoid_ilt_mask(mask1)
#     # 打印mask1的最大最小值
#     print(f"mask1最大值: {mask1.max().item()}, mask1最小值: {mask1.min().item()}")
#     theta1_back = torch.acos(2*mask1 - 1)
#     theta1 = theta1_back.clone().detach().requires_grad_(True)
#     # 动态调整学习率
#     optimizer1 = optim.Adam([theta1], lr=0.1)
#     # optimizer1 = optim.Adam([theta1], lr=0.01, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-4)
#
#     # 使用sigmod松弛
#     # p = torch.logit(mask1.clamp(1e-6, 1 - 1e-6))
#     # p = p.clone().detach().requires_grad_(True)
#     # optimizer1 = optim.Adam([p], lr=0.1)
#
#     # 自适应学习率
#     scheduler1 = optim.lr_scheduler.ReduceLROnPlateau(optimizer1, mode='min', factor=0.5, patience=5, verbose=True)
#     # optimizer1 = optim.SGD([theta1], lr=0.1, momentum=0.9, dampening=0, weight_decay=1e-3,
#     #                     nesterov=True)
#     # 定义热重启学习率策略
#     # scheduler1 = lr_scheduler.CosineAnnealingWarmRestarts(optimizer1, T_0=10, T_mult=2, eta_min=0, last_epoch=-1)
#
#     best_loss1 = float('inf')
#     best_mask1 = None
#     best_intensity_map1 = None
#     best_binary_wafer1 = None
#
#     target1_lf = ut.ideal_lpfilter(target1, 193, 1.35)
#
#     grad_values = []
#
#     initial_weight = 100
#     decay_rate = 0.99
#
#     # Run optimization
#     print(f'缩小4倍，分辨率为 {mask1.shape[2]} x {mask1.shape[3]}: ')
#     for epoch in range(200):
#         mask1_tmp = (torch.cos(theta1) + 1) / 2
#
#         regu1 = torch.sum(mask1_tmp ** 2 * (mask1_tmp - 1) ** 2)
#         regu1_weight = initial_weight * (decay_rate ** epoch)  # 指数衰减
#         # 使用 sigmoid 将 p 映射到 [0,1] 得到 mask
#         # mask1_tmp = torch.sigmoid(p)
#
#         mask1_tmp_lf = ut.ideal_lpfilter(mask1_tmp, 193, 1.35)
#         # l2_loss1 = ilt_loss.apply(mask1_tmp_lf, target1_lf, kernel_focus_lowres4, kernel_focus_ct_lowres4, weight_focus)
#         l2_loss1 = ilt_loss.apply(mask1_tmp_lf, target1_lf, kernel_focus, kernel_focus_ct, weight_focus)
#         # print(f"l2_loss1: {l2_loss1}")
#         # print(f"regu1: {regu1}")
#         # inten_nom, binary_nom_lf = litho.lithosim(mask1_tmp_lf, 0.225, kernel_focus_lowres4, weight_focus,
#         #                             wafer_output_path=None, save_bin_wafer_image=False, return_binary_wafer=True)
#         # l2_loss1 = func.mse_loss(binary_nom_lf, target1_lf, reduction="sum")
#         loss1 = l2_loss1 + regu1_weight * regu1
#         # loss1 = l2_loss1
#         # intensity1,binary1 = litho.lithosim(mask1, 0.225, kernel_focus_lowres4, weight_focus, None, False)
#         # loss1 = func.mse_loss(binary1, target1, reduction="sum")
#         # print(mask1.requires_grad)  # 这应该返回 True
#         # print(loss1.grad_fn)  # 如果是 None，说明计算图断裂
#         print(f'缩小四倍 Epoch {epoch + 1}, Loss: {loss1.item():.4f}')
#
#         # 记录最佳结果
#         if best_loss1 is None or best_mask1 is None or loss1.item() < best_loss1:
#             best_loss1 = loss1.detach().clone()
#             best_mask1 = mask1_tmp.detach().clone()
#
#         optimizer1.zero_grad()
#         loss1.backward()
#
#         # # 打印梯度信息
#         # print(f"Epoch {epoch + 1}: Gradient Statistics")
#         # for name, param in [('theta1', theta1)]:
#         #     if param.grad is not None:
#         #         print(f"  {name}.grad - max: {param.grad.max().item():.6f}, "
#         #               f"min: {param.grad.min().item():.6f}, "
#         #               f"mean: {param.grad.mean().item():.6f}")
#         #     else:
#         #         print(f"  {name}.grad is None!")  # 可能计算图断裂
#
#         # 记录梯度均值
#         grad_values.append(theta1.grad.abs().mean().item())
#
#         optimizer1.step()
#
#         # 调度器根据当前 loss 调整学习率
#         scheduler1.step(loss1.item())
#
#     # 绘制梯度变化趋势
#     plt.plot(grad_values)
#     plt.xlabel('Epoch')
#     plt.ylabel('Gradient Mean')
#     plt.title('Gradient Vanishing/Exploding Check')
#     plt.show()
#
#     # Save the best mask image
#     if best_mask1 is not None:
#         best_intensity_map1, best_binary_wafer1 = litho.lithosim(best_mask1, 0.225, kernel_focus, weight_focus, None,
#                                                                False)
#
#         plt.imshow(best_mask1.squeeze().cpu().numpy())
#         plt.title('4倍降采样掩模优化结果')
#         plt.colorbar()
#         plt.show()
#         # best_mask_image = Image.fromarray(best_mask_image.astype('float32'))
#         # best_mask_image.save('best_mask1.png')
#         # print('Best mask saved as best_mask1.png')
#
#     # Save the best intensity_map and binary_wafer
#     if best_intensity_map1 is not None and best_binary_wafer1 is not None:
#         # best_intensity_map_image = best_intensity_map1.squeeze().cpu().numpy()
#         # best_intensity_map_image = Image.fromarray(best_intensity_map_image.astype('uint8'))
#         # best_intensity_map_image.save('best_intensity_map2.png')
#         # print('Best intensity map saved as best_intensity_map2.png')
#         # #
#         # best_binary_wafer_image = best_binary_wafer1.squeeze().cpu().numpy()
#         # best_binary_wafer_image = Image.fromarray(best_binary_wafer_image.astype('uint8'))
#         # best_binary_wafer_image.save('best_binary_wafer2.png')
#         # print('Best binary wafer saved as best_binary_wafer2.png')
#
#         plt.imshow(best_intensity_map1.squeeze().cpu().numpy())
#         plt.title('4倍降采样空间像')
#         plt.colorbar()
#         plt.show()
#
#         plt.imshow(best_binary_wafer1.squeeze().cpu().numpy())
#         plt.title('4倍降采样光刻胶像')
#         plt.colorbar()
#         plt.show()
#
#     ####################################################################################################################
#     # 2.将best_mask上采样作为输入，继续做梯度下降
#     print("#################################################################################")
#     target2 = target_lowres2
#     target2 = target2.clone().detach().requires_grad_(False)
#
#     mask_tmp = torch.nn.functional.interpolate(best_mask1, size=(1024, 1024), mode='bilinear')
#     mask2 = mask_tmp
#     mask2 = mask2.clone().detach().requires_grad_(False)
#     theta2_back = torch.acos(2 * mask2 - 1)
#     theta2 = theta2_back.clone().detach().requires_grad_(True)
#     optimizer2 = optim.Adam([theta2], lr=0.01)
#     # 自适应学习率
#     scheduler2 = optim.lr_scheduler.ReduceLROnPlateau(optimizer2, mode='min', factor=0.5, patience=5, verbose=True)
#
#     best_loss2 = float('inf')
#     best_mask2 = None
#     best_intensity_map2 = None
#     best_binary_wafer2 = None
#
#     target2_lf = ut.ideal_lpfilter(target2, 193, 1.35)
#
#     # 优化循环
#     print(f'缩小2倍，分辨率为 {mask2.shape[2]} x {mask2.shape[3]}: ')
#     for epoch in range(100):
#         mask2_tmp = (torch.cos(theta2) + 1) / 2
#         mask2_tmp_lf = ut.ideal_lpfilter(mask2_tmp, 193, 1.35)
#         l2_loss2 = ilt_loss.apply(mask2_tmp_lf, target2_lf, kernel_focus, kernel_focus_ct, weight_focus)
#         loss2 = l2_loss2
#
#         print(f'缩小2倍 Epoch {epoch + 1}, Loss: {loss2.item():.4f}')
#
#         # 记录最佳结果
#         if best_loss2 is None or best_mask2 is None or loss2.item() < best_loss2:
#             best_loss2 = loss2.detach().clone()
#             best_mask2 = mask2_tmp.detach().clone()
#
#         optimizer2.zero_grad()
#         loss2.backward()
#
#         # # 记录梯度均值
#         # grad_values.append(theta2.grad.abs().mean().item())
#         optimizer2.step()
#         # 调度器根据当前 loss 调整学习率
#         scheduler1.step(loss1.item())
#
#     # # 绘制梯度变化趋势
#     # plt.plot(grad_values)
#     # plt.xlabel('Epoch')
#     # plt.ylabel('Gradient Mean')
#     # plt.title('Gradient Vanishing/Exploding Check')
#     # plt.show()
#
#     # Save the best mask image
#     if best_mask2 is not None:
#         best_intensity_map2, best_binary_wafer2 = litho.lithosim(best_mask2, 0.225, kernel_focus, weight_focus,
#                                                                  None,False)
#         best_mask_image = best_mask2.squeeze().cpu().numpy()
#         plt.imshow(best_mask_image)
#         plt.title('2倍降采样掩模优化结果')
#         plt.show()
#         # best_mask_image = Image.fromarray(best_mask_image.astype('float32'))
#         # best_mask_image.save('best_mask2.png')
#         # print('Best mask saved as best_mask2.png')
#
#     # Save the best intensity_map and binary_wafer
#     if best_intensity_map2 is not None and best_binary_wafer2 is not None:
#         # best_intensity_map_image = best_intensity_map2.squeeze().cpu().numpy()
#         # best_intensity_map_image = Image.fromarray(best_intensity_map_image.astype('uint8'))
#         # best_intensity_map_image.save('best_intensity_map2.png')
#         # print('Best intensity map saved as best_intensity_map2.png')
#         # #
#         # best_binary_wafer_image = best_binary_wafer2.squeeze().cpu().numpy() * 255
#         # best_binary_wafer_image = Image.fromarray(best_binary_wafer_image.astype('uint8'))
#         # best_binary_wafer_image.save('best_binary_wafer2.png')
#         # print('Best binary wafer saved as best_binary_wafer2.png')
#
#         plt.imshow(best_intensity_map2.squeeze().cpu().numpy())
#         plt.title('2倍降采样空间像')
#         plt.show()
#         plt.imshow(best_binary_wafer2.squeeze().cpu().numpy())
#         plt.title('2倍降采样光刻胶像')
#         plt.show()
#
#     ####################################################################################################################
#     # 3.插值回原分辨率
#     print("#################################################################################")
#     target3 = target_highres
#     target3 = target3.clone().detach().requires_grad_(False)
#     mask_tmp2 = torch.nn.functional.interpolate(best_mask2, size=(2048, 2048), mode='bilinear')
#     mask3 = mask_tmp2
#     mask3 = mask3.clone().detach().requires_grad_(False)
#     theta3_back = torch.acos(2 * mask3 - 1)
#     theta3 = theta3_back.clone().detach().requires_grad_(True)
#     optimizer3 = optim.Adam([theta3], lr=0.01)
#     target3_lf = ut.ideal_lpfilter(target3, 193, 1.35)
#     # 自适应学习率
#     scheduler3 = optim.lr_scheduler.ReduceLROnPlateau(optimizer3, mode='min', factor=0.5, patience=5, verbose=True)
#
#
#     best_loss3 = float('inf')
#     best_mask3 = None
#     best_intensity_map3 = None
#     best_binary_wafer3 = None
#
#     # 优化循环
#     print(f'原尺寸 分辨率为 {mask3.shape[2]} x {mask3.shape[3]}: ')
#     for epoch in range(10):
#         mask3_tmp = (torch.cos(theta3) + 1) / 2
#         mask3_tmp_lf = ut.ideal_lpfilter(mask3_tmp, 193, 1.35)
#         l2_loss3 = ilt_loss.apply(mask3_tmp_lf, target3_lf, kernel_focus, kernel_focus_ct, weight_focus)
#         loss3 = l2_loss3
#
#         print(f'原尺寸 Epoch {epoch + 1}, Loss: {loss3.item():.4f}')
#
#         # 记录最佳结果
#         if best_loss3 is None or best_mask3 is None or loss3.item() < best_loss3:
#             best_loss3 = loss3.detach().clone()
#             best_mask3 = mask3_tmp.detach().clone()
#             # plt.imshow(best_mask3.squeeze().cpu().numpy())
#             # plt.show()
#
#         optimizer3.zero_grad()
#         loss3.backward()
#         optimizer3.step()
#
#
#     # Save the best mask image
#     if best_mask3 is not None:
#         best_intensity_map3, best_binary_wafer3 = litho.lithosim(best_mask3, 0.225, kernel_focus, weight_focus, None,
#                                                                False)
#         best_mask_image = best_mask3.squeeze().cpu().numpy()
#         plt.imshow(best_mask_image)
#         plt.title('原尺寸掩模优化结果')
#         plt.show()
#         # best_mask_image = Image.fromarray(best_mask_image.astype('float32'))
#         # best_mask_image.save('best_mask3.png')
#
#     # Save the best intensity_map and binary_wafer
#     if best_intensity_map3 is not None and best_binary_wafer3 is not None:
#         # best_intensity_map_image = best_intensity_map3.squeeze().cpu().numpy()
#         # best_intensity_map_image = Image.fromarray(best_intensity_map_image.astype('uint8'))
#         # best_intensity_map_image.save('best_intensity_map3.png')
#         # print('Best intensity map saved as best_intensity_map3.png')
#         # #
#         # best_binary_wafer_image = best_binary_wafer3.squeeze().cpu().numpy() * 255
#         # best_binary_wafer_image = Image.fromarray(best_binary_wafer_image.astype('uint8'))
#         # best_binary_wafer_image.save('best_binary_wafer3.png')
#         # print('Best binary wafer saved as best_binary_wafer3.png')
#         plt.imshow(best_intensity_map3.squeeze().cpu().numpy())
#         plt.title('原尺寸空间像')
#         plt.show()
#         plt.imshow(best_binary_wafer3.squeeze().cpu().numpy())
#         plt.title('原尺寸光刻胶像')
#         plt.show()



# 2025-1-27版本，梯度滤波和kerenl未降采样
# if __name__ == '__main__':
#     # 基础配置
#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#
#     # 加载预训练核
#     kernel_focus = torch.load('kernel_neuralilt/kernel_focus_tensor.pt', map_location=device)
#     kernel_focus_ct = torch.load('kernel_neuralilt/kernel_ct_focus_tensor.pt', map_location=device)
#     kernel_defocus = torch.load('kernel_neuralilt/kernel_defocus_tensor.pt', map_location=device)
#     kernel_defocus_ct = torch.load('kernel_neuralilt/kernel_ct_defocus_tensor.pt', map_location=device)
#     weight_focus = torch.load('kernel_neuralilt/weight_focus_tensor.pt', map_location=device)
#     weight_defocus = torch.load('kernel_neuralilt/weight_defocus_tensor.pt', map_location=device)
#
#     mask_path = 'ICCAD2013/png/target1.png'
#     mask_highres = ut.load_image(mask_path)  # [1,1,2048,2048]
#     mask_highres = ut.mask_init(mask_highres)  # [1,1,2048,2048]
#
#     # mask_highres = ut.ideal_lpfilter(mask_highres, lamda=193, NA=0.8)  # [1,1,2048,2048]
#
#     target_highres = ut.load_image(mask_path)
#
#     # 降采样
#     avgpool_size = 4  # 降采样因子
#     avg_layer = nn.AvgPool2d(avgpool_size, stride=avgpool_size)
#     mask_lowres = avg_layer(mask_highres)  # [1,512,512]
#     target_lowres = avg_layer(target_highres)  # [1,512,512]
#     # plt.imshow(mask_lowres.squeeze().cpu().numpy())
#     # plt.show()
#
#     avgpool_size2 = 2  # 降采样因子
#     avg_layer2 = nn.AvgPool2d(avgpool_size2, stride=avgpool_size2)
#     target_lowres2 = avg_layer2(target_highres)
#     # plt.imshow(target_lowres2.squeeze().cpu().numpy())
#     # plt.show()
#
#     # 1.先降采样，然后做梯度下降
#     mask1 = mask_lowres
#     target1 = target_lowres
#
#     # Ensure mask has the correct dimensions
#     if len(mask1.shape) == 3:
#         mask1 = mask1.unsqueeze(0)  # Add batch dimension
#
#     # Define optimizer
#     mask1.requires_grad_(True)
#     optimizer1 = optim.Adam([mask1], lr=0.01)
#
#     best_loss1 = float('inf')
#     best_mask1 = None
#     best_intensity_map1 = None
#     best_binary_wafer1 = None
#
#     # Run optimization
#     print(f'缩小4倍，分辨率为 {mask1.shape[2]} x {mask1.shape[3]}: ')
#     for epoch in range(60):
#         optimizer1.zero_grad()
#
#         loss1 = ilt_loss.apply(mask1, target1, kernel_focus, kernel_focus_ct, weight_focus)
#         loss1.backward()
#
#         # 梯度滤波
#         with torch.no_grad():
#             for param in optimizer1.param_groups[0]['params']:
#                 if param.grad is not None:
#                     # 应用低通滤波且保留原始幅度
#                     raw_grad_norm = torch.norm(param.grad)
#                     filtered_grad = ut.ideal_lpfilter(param.grad, lamda=193, NA=0.8)
#
#                     # 保持原始梯度幅度
#                     filtered_grad *= raw_grad_norm / (torch.norm(filtered_grad) + 1e-8)
#
#                     param.grad = filtered_grad
#
#         # 参数更新
#         optimizer1.step()
#
#         # 记录最佳结果
#         if loss1.item() < best_loss1:
#             best_loss1 = loss1.item()
#             best_mask1 = mask1.clone().detach()
#
#         print(f'缩小四倍 Epoch {epoch + 1}, Loss: {loss1.item():.4f}')
#
#     # Save the best mask image
#     if best_mask1 is not None:
#         best_intensity_map1, best_binary_wafer1 = litho.lithosim(best_mask1, 0.225, kernel_focus, weight_focus, None,
#                                                                False)
#         best_mask_image = best_mask1.squeeze().cpu().numpy()
#         plt.imshow(best_mask_image)
#         plt.show()
#         # best_mask_image = Image.fromarray(best_mask_image.astype('float32'))
#         # best_mask_image.save('best_mask1.png')
#         print('Best mask saved as best_mask1.png')
#
#     # Save the best intensity_map and binary_wafer
#     if best_intensity_map1 is not None and best_binary_wafer1 is not None:
#         best_intensity_map_image = best_intensity_map1.squeeze().cpu().numpy()
#         best_intensity_map_image = (best_intensity_map_image - best_intensity_map_image.min()) / (
#                 best_intensity_map_image.max() - best_intensity_map_image.min()) * 255
#         best_intensity_map_image = Image.fromarray(best_intensity_map_image.astype('uint8'))
#         best_intensity_map_image.save('best_intensity_map1.png')
#         print('Best intensity map saved as best_intensity_map1.png')
#
#         best_binary_wafer_image = best_binary_wafer1.squeeze().cpu().numpy() * 255
#         best_binary_wafer_image = Image.fromarray(best_binary_wafer_image.astype('uint8'))
#         best_binary_wafer_image.save('best_binary_wafer1.png')
#         print('Best binary wafer saved as best_binary_wafer1.png')
#
#     ####################################################################################################################
#     # 2.将best_mask上采样作为输入，继续做梯度下降
#     mask_tmp = torch.nn.functional.interpolate(best_mask1, size=(1024, 1024), mode='nearest')
#     best_loss2 = float('inf')
#     best_mask2 = None
#     best_intensity_map2 = None
#     best_binary_wafer2 = None
#
#     # # 使用matplotlib绘制best_mask和mask_tmp
#     # plt.imshow(best_mask.squeeze().cpu().numpy())
#     # plt.show()
#     # plt.imshow(mask_tmp.squeeze().cpu().numpy())
#     # plt.show()
#
#     mask2 = mask_tmp.clone().detach()
#     # mask2 = ut.mask_init(mask2)
#     target2 = target_lowres2
#     print(f'缩小2倍，分辨率为 {mask2.shape[2]} x {mask2.shape[3]}: ')
#
#     mask2.requires_grad_(True)
#     optimizer2 = optim.Adam([mask2], lr=0.01)
#     # 优化循环
#     for epoch in range(200):
#         optimizer2.zero_grad()
#
#         # 前向传播
#         loss2 = ilt_loss.apply(mask2, target2, kernel_focus, kernel_focus_ct, weight_focus)
#
#         # 反向传播
#         loss2.backward()
#
#         # 梯度滤波 (关键修改点)
#         with torch.no_grad():
#             for param in optimizer2.param_groups[0]['params']:
#                 if param.grad is not None:
#                     # 应用低通滤波且保留原始幅度
#                     raw_grad_norm = torch.norm(param.grad)
#                     filtered_grad = ut.ideal_lpfilter(param.grad, lamda=193, NA=0.8)
#
#                     # 保持原始梯度幅度
#                     filtered_grad *= raw_grad_norm / (torch.norm(filtered_grad) + 1e-8)
#
#                     param.grad = filtered_grad
#
#         # 参数更新
#         optimizer2.step()
#
#         # 记录最佳结果
#         if loss2.item() < best_loss2:
#             best_loss2 = loss2.item()
#             best_mask2 = mask2.clone().detach()
#
#         print(f'缩小2倍 Epoch {epoch + 1}, Loss: {loss2.item():.4f}')
#
#     # Save the best mask image
#     if best_mask2 is not None:
#         best_intensity_map2, best_binary_wafer2 = litho.lithosim(best_mask2, 0.225, kernel_focus, weight_focus, None,
#                                                                False)
#         best_mask_image = best_mask2.squeeze().cpu().numpy()
#         plt.imshow(best_mask_image)
#         plt.show()
#         # best_mask_image = Image.fromarray(best_mask_image.astype('float32'))
#         # best_mask_image.save('best_mask2.png')
#         # print('Best mask saved as best_mask2.png')
#
#     # Save the best intensity_map and binary_wafer
#     if best_intensity_map2 is not None and best_binary_wafer2 is not None:
#         best_intensity_map_image = best_intensity_map2.squeeze().cpu().numpy()
#         best_intensity_map_image = (best_intensity_map_image - best_intensity_map_image.min()) / (
#                 best_intensity_map_image.max() - best_intensity_map_image.min()) * 255
#         # best_intensity_map_image = 255 - best_intensity_map_image  # Invert colors
#         best_intensity_map_image = Image.fromarray(best_intensity_map_image.astype('uint8'))
#         best_intensity_map_image.save('best_intensity_map2.png')
#         print('Best intensity map saved as best_intensity_map2.png')
#
#         best_binary_wafer_image = best_binary_wafer2.squeeze().cpu().numpy() * 255
#         # best_binary_wafer_image = 255 - best_binary_wafer_image  # Invert colors
#         best_binary_wafer_image = Image.fromarray(best_binary_wafer_image.astype('uint8'))
#         best_binary_wafer_image.save('best_binary_wafer2.png')
#         print('Best binary wafer saved as best_binary_wafer2.png')
#
#     ####################################################################################################################
#     # 3.插值回原分辨率
#     mask_tmp2 = torch.nn.functional.interpolate(best_mask2, size=(2048, 2048), mode='nearest')
#     mask3 = mask_tmp2.clone().detach()
#     # mask3 = ut.mask_init(mask3)
#     target3 = target_highres
#     print(f'原尺寸 分辨率为 {mask3.shape[2]} x {mask3.shape[3]}: ')
#     mask3.requires_grad_(True)
#     optimizer3 = optim.Adam([mask3], lr=0.01)
#
#     best_loss3 = float('inf')
#     best_mask3 = None
#     best_intensity_map3 = None
#     best_binary_wafer3 = None
#
#     for epoch in range(70):
#         optimizer3.zero_grad()
#         loss3 = ilt_loss.apply(mask3, target3, kernel_focus, kernel_focus_ct, weight_focus)
#         loss3.backward()
#
#         # 梯度滤波
#         with torch.no_grad():
#             for param in optimizer3.param_groups[0]['params']:
#                 if param.grad is not None:
#                     # 应用低通滤波且保留原始幅度
#                     raw_grad_norm = torch.norm(param.grad)
#                     filtered_grad = ut.ideal_lpfilter(param.grad, lamda=193, NA=0.8)
#
#                     # 保持原始梯度幅度
#                     filtered_grad *= raw_grad_norm / (torch.norm(filtered_grad) + 1e-8)
#                     param.grad = filtered_grad
#
#         # 参数更新
#         optimizer3.step()
#
#         # 记录最佳结果
#         if loss3.item() < best_loss3:
#             best_loss3 = loss3.item()
#             best_mask3 = mask3.clone().detach()
#             plt.imshow(best_mask3.squeeze().cpu().numpy())
#             plt.show()
#
#         print(f'原尺寸 Epoch {epoch + 1}, Loss: {loss3.item():.4f}')
#
#     # Save the best mask image
#     if best_mask3 is not None:
#         best_intensity_map3, best_binary_wafer3 = litho.lithosim(best_mask3, 0.225, kernel_focus, weight_focus, None,
#                                                                False)
#         best_mask_image = best_mask3.squeeze().cpu().numpy()
#         # best_mask_image = (best_mask_image + 1) / 2 * 255  # Convert to [0, 255] range
#         # best_mask = best_mask_image * 255
#         # best_mask_image = 255 - best_mask_image  # Invert colors
#         plt.imshow(best_mask_image)
#         plt.show()
#         # best_mask_image = Image.fromarray(best_mask_image.astype('float32'))
#         # best_mask_image.save('best_mask3.png')
#         print('Best mask saved as best_mask3.png')
#
#     # Save the best intensity_map and binary_wafer
#     if best_intensity_map3 is not None and best_binary_wafer3 is not None:
#         best_intensity_map_image = best_intensity_map3.squeeze().cpu().numpy()
#         best_intensity_map_image = (best_intensity_map_image - best_intensity_map_image.min()) / (
#                 best_intensity_map_image.max() - best_intensity_map_image.min()) * 255
#         best_intensity_map_image = 255 - best_intensity_map_image  # Invert colors
#         best_intensity_map_image = Image.fromarray(best_intensity_map_image.astype('uint8'))
#         best_intensity_map_image.save('best_intensity_map3.png')
#         print('Best intensity map saved as best_intensity_map3.png')
#
#         best_binary_wafer_image = best_binary_wafer3.squeeze().cpu().numpy() * 255
#         # best_binary_wafer_image = 255 - best_binary_wafer_image  # Invert colors
#         best_binary_wafer_image = Image.fromarray(best_binary_wafer_image.astype('uint8'))
#         best_binary_wafer_image.save('best_binary_wafer3.png')
#         print('Best binary wafer saved as best_binary_wafer3.png')


# 未降采样版本
# if __name__ == '__main__':
#     # 基础配置
#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#     # original_size = 2048
#
#     # 加载预训练核
#     kernel_focus = torch.load('kernel_neuralilt/kernel_focus_tensor.pt', map_location=device)
#     kernel_focus_ct = torch.load('kernel_neuralilt/kernel_ct_focus_tensor.pt', map_location=device)
#     kernel_defocus = torch.load('kernel_neuralilt/kernel_defocus_tensor.pt', map_location=device)
#     kernel_defocus_ct = torch.load('kernel_neuralilt/kernel_ct_defocus_tensor.pt', map_location=device)
#     weight_focus = torch.load('kernel_neuralilt/weight_focus_tensor.pt', map_location=device)
#     weight_defocus = torch.load('kernel_neuralilt/weight_defocus_tensor.pt', map_location=device)
#
#     mask_path = 'ICCAD2013/png/target1.png'
#     mask = ut.load_image(mask_path)
#
#     target = ut.load_image(mask_path)
#
#     # 初始化掩模
#     mask = ut.mask_init(mask)
#     lambda_ = 193
#     NA = 0.8
#     mask = ut.ideal_lpfilter(mask, lambda_, NA)
#     # Ensure mask has the correct dimensions
#     if len(mask.shape) == 3:
#         mask = mask.unsqueeze(0)  # Add batch dimension
#
#     # 定义损失函数
#     ilt_loss_fn = ilt_loss.apply
#     # ilt_loss_fn = ilt_loss_scale_function.apply
#
#     # Define optimizer
#     mask.requires_grad_(True)
#     optimizer = optim.Adam([mask], lr=0.01)
#
#     best_loss = float('inf')
#     best_mask = None
#     best_intensity_map = None
#     best_binary_wafer = None
#
#     # Run optimization
#     for epoch in range(200):
#         optimizer.zero_grad()
#         loss = ilt_loss_fn(mask, target, kernel_focus, kernel_focus_ct, weight_focus)
#         loss.backward()
#         optimizer.step()
#         print(f'Epoch {epoch + 1}, Loss: {loss.item()}')
#
#         # Save the best mask and corresponding intensity_map, binary_wafer
#         if loss.item() < best_loss:
#             best_loss = loss.item()
#             best_mask = mask.clone().detach()
#             best_intensity_map, best_binary_wafer = litho.lithosim(best_mask, 0.225, kernel_focus, weight_focus, None,
#                                                                    False)
#
#     # Save the best mask image
#     if best_mask is not None:
#         best_mask_image = best_mask.squeeze().cpu().numpy()
#         best_mask_image = (best_mask_image + 1) / 2 * 255  # Convert to [0, 255] range
#         # best_mask_image = 255 - best_mask_image  # Invert colors
#         best_mask_image = Image.fromarray(best_mask_image.astype('uint8'))
#         best_mask_image.save('best_mask.png')
#         print('Best mask saved as best_mask.png')
#
#     # Save the best intensity_map and binary_wafer
#     if best_intensity_map is not None and best_binary_wafer is not None:
#         best_intensity_map_image = best_intensity_map.squeeze().cpu().numpy()
#         best_intensity_map_image = (best_intensity_map_image - best_intensity_map_image.min()) / (
#                 best_intensity_map_image.max() - best_intensity_map_image.min()) * 255
#         # best_intensity_map_image = 255 - best_intensity_map_image  # Invert colors
#         best_intensity_map_image = Image.fromarray(best_intensity_map_image.astype('uint8'))
#         best_intensity_map_image.save('best_intensity_map.png')
#         print('Best intensity map saved as best_intensity_map.png')
#
#         best_binary_wafer_image = best_binary_wafer.squeeze().cpu().numpy() * 255
#         # best_binary_wafer_image = 255 - best_binary_wafer_image  # Invert colors
#         best_binary_wafer_image = Image.fromarray(best_binary_wafer_image.astype('uint8'))
#         best_binary_wafer_image.save('best_binary_wafer.png')
#         print('Best binary wafer saved as best_binary_wafer.png')



# 降采样版本
# if __name__ == '__main__':
#     # 基础配置
#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#     original_size = 2048  # 原始掩模尺寸
#     downsampled_size = 512  # 降采样后尺寸（论文中avgpool_size=4）
#     avgpool_size = 4  # 降采样因子
#
#     # 加载SOCS核（需确保核尺寸适配降采样后分辨率）
#     # Load pre-trained kernels
#     kernel_focus = torch.load('kernel_neuralilt/kernel_focus_tensor.pt', map_location=device)
#     kernel_focus_ct = torch.load('kernel_neuralilt/kernel_ct_focus_tensor.pt', map_location=device)
#     kernel_defocus = torch.load('kernel_neuralilt/kernel_defocus_tensor.pt', map_location=device)
#     kernel_defocus_ct = torch.load('kernel_neuralilt/kernel_ct_defocus_tensor.pt', map_location=device)
#     weight_focus = torch.load('kernel_neuralilt/weight_focus_tensor.pt', map_location=device)
#     weight_defocus = torch.load('kernel_neuralilt/weight_defocus_tensor.pt', map_location=device)
#
#     # 加载目标布局并预处理
#     mask_path = 'ICCAD2013/png/target1.png'
#     mask_highres = ut.load_image(mask_path)  # [1,1,2048,2048]
#     target_highres = ut.load_image(mask_path)
#
#     # Step 1: 目标布局降采样（匹配仿真输出尺寸）
#     avg_layer = nn.AvgPool2d(avgpool_size, stride=avgpool_size)
#     target_lowres = avg_layer(target_highres)  # [1,512,512]
#     target_lowres = target_lowres.unsqueeze(0)  # [1,1,512,512]
#
#
#     # Step 2: 初始化可优化掩模（低分辨率）
#     mask_highres = ut.mask_init(mask_highres)  # [1,1,2048,2048]
#     mask_lowres = avg_layer(mask_highres)  # [1,512,512]
#     mask_lowres = mask_lowres.unsqueeze(0)  # [1,1,512,512]
#
#     # Step 3: 定义裁剪坐标（new_cord参数）
#     # 假设优化区域为中心512x512区域（需根据实际布局调整）
#     lx = (original_size - downsampled_size) // 2
#     ly = lx
#     rx = lx + downsampled_size
#     ry = ly + downsampled_size
#     new_cord = (
#         torch.tensor([lx], device=device),
#         torch.tensor([ly], device=device),
#         torch.tensor([rx], device=device),
#         torch.tensor([ry], device=device)
#     )
#
#     # Step 4: 配置优化器
#     mask_lowres.requires_grad_(True)
#     optimizer = optim.Adam([mask_lowres], lr=0.01)
#     best_loss = float('inf')
#     best_mask = None
#
#     # 优化循环
#     for epoch in range(200):
#         optimizer.zero_grad()
#
#         # 前向传播（自动处理缩放）
#         loss, _ = ilt_loss_scale_function.apply(
#             mask_lowres,
#             target_lowres,
#             kernel_focus,
#             kernel_focus_ct,
#             kernel_defocus,  # kernel_def（未使用复杂度优化）
#             kernel_defocus_ct,  # kernel_def_ct
#             weight_focus,
#             weight_defocus,  # weight_def
#             new_cord,
#             False,
#             False,
#             False
#         )
#
#         # 反向传播
#         loss.backward()
#         optimizer.step()
#
#         # 记录最佳结果
#         if loss.item() < best_loss:
#             best_loss = loss.item()
#             best_mask = mask_lowres.detach().clone()
#
#         print(f'Epoch {epoch + 1}, Loss: {loss.item():.4f}')
#
#     # 后处理：将优化后的低分辨率掩模映射回原始尺寸
#     if best_mask is not None:
#         # Step 5: 上采样掩模（最近邻插值保持二值特性）
#         mask_upsampled = torch.nn.functional.interpolate(
#             best_mask,
#             size=(original_size, original_size),
#             mode='nearest'
#         )
#
#         # Step 6: 保存优化结果
#         mask_np = (mask_upsampled.squeeze().cpu().numpy() > 0).astype(np.uint8) * 255
#         Image.fromarray(mask_np).save('optimized_mask_highres.png')
#
#         # 可选：生成高分辨率光刻仿真结果
#         with torch.no_grad():
#             intensity, wafer = litho.lithosim(
#                 mask_upsampled,
#                 threshold=0.225,
#                 kernels=kernel_focus,  # 需使用原始分辨率核
#                 weight=weight_focus,
#                 wafer_output_path=None,
#                 save_bin_wafer_image=False,
#                 avgpool_size=None  # 禁用降采样
#             )
#             # 保存光强图
#             intensity_np = intensity.squeeze().cpu().numpy()
#             intensity_norm = (intensity_np - intensity_np.min()) / (intensity_np.max() - intensity_np.min()) * 255
#             Image.fromarray(intensity_norm.astype(np.uint8)).save('intensity_highres.png')
#             wafer_np = wafer.squeeze().cpu().numpy() * 255
#             Image.fromarray(wafer_np.astype(np.uint8)).save('wafer_highres.png')


# class ilt_loss_scale_function(Function):
#     r"""
#     A wrapper class of ilt_loss and cplx_loss calculations with scaling scheme
#     Scaling scheme:
#         Input mask (nn prediction) -> scale back to the cropped bbox size -> fit the cropped bbox into original size
#                                    -> do litho-simulation -> compute ilt_loss/cplx_loss
#                                    -> calculate the gradient of ilt_loss/cplx_loss -> crop and scale the gradient to input size -> gradient backward to nn
#         [input_size, input_size] -> [cropped_bbox_size, cropped_bbox_size] -> [original_size, original_size] -> [cropped_bbox_size, cropped_bbox_size] -> [input_size, input_size]
#     """
#
#     @staticmethod
#     def forward(ctx, mask_pred, target, kernels, kernels_ct, kernel_def, kernel_def_ct, weight, weight_def, new_cord,
#                 cycle_mode=False, cplx_obj=False, report_epe=False):
#         lx, ly, rx, ry = new_cord
#         output_litho_l2_loss = True
#
#         # The exact ILT forward loss
#         # ilt_loss = ||litho(threshold(Phi(Z_t, w)), P_nom) - Z_t||_gamma
#         mask_pred_sig = ut.mask_init(mask_pred) # change mask_pred \in {0,1} to \in {-1,1}
#         mask_pred_sig = ut.sigmoid_ilt_mask(mask_pred_sig, theta_m=4)
#         mask_pred_sig_backup = torch.clone(mask_pred_sig)
#         if output_litho_l2_loss:
#             # It is easier to monitor thru L2 loss, rather than the exact ILT loss
#             mask_pred_sig = (mask_pred > 0.5).type(mask_pred.dtype)
#
#         batch_size = mask_pred_sig.shape[0]
#         channel_size = mask_pred_sig.shape[1]
#
#         mask_pred_sig_orig_size = torch.zeros((batch_size, channel_size, 2048, 2048),
#                                               dtype=mask_pred_sig.dtype, layout=mask_pred_sig.layout,
#                                               device=mask_pred_sig.device)
#         for i in range(batch_size):
#             cur_mask = mask_pred_sig[i].unsqueeze(0)  # 1 * 1 * H * W
#
#             mask_crop = torch.nn.functional.interpolate(cur_mask,
#                                                         size=(abs(rx - lx)[i].item(), abs(rx - lx)[i].item()),
#                                                         mode='nearest')  # 1 * 1 * H * W
#             mask_origin = torch.zeros((mask_crop.shape[0], channel_size, 2048, 2048),
#                                       dtype=mask_crop.dtype, layout=mask_crop.layout, device=mask_crop.device)
#             mask_origin[..., ly[i].item():ry[i].item(), lx[i].item():rx[i].item()] = mask_crop
#             mask_pred_sig_orig_size[i] = mask_origin.squeeze(0)
#
#         # avgpool_size = 4
#         mask_pred_bin_orig_size = (mask_pred_sig_orig_size > 0.5).type(torch.cuda.FloatTensor)
#
#         # result, bin_mask = litho.lithosim(mask_pred_bin_orig_size, 0.225, kernels, weight,
#         #                                   None, False, return_binary_wafer=True)
#         result, bin_mask = litho.lithosim(mask_pred_sig, 0.225, kernels, weight,
#                                           None, False, return_binary_wafer=True)
#
#         if cycle_mode:
#             # It's easier to monitor the L2 loss using mean function but not the sum
#             ilt_loss = (result - target).pow(4).mean()
#             l2_loss = (bin_mask - target).abs().mean()
#         elif cplx_obj:
#             result_inner, bin_mask_inner = litho.lithosim(mask_pred_bin_orig_size, 0.225, kernel_def, weight_def,
#                                                           None, False, avgpool_size=avgpool_size,
#                                                           return_binary_wafer=True, dose=0.98)
#             result_outer, bin_mask_outer = litho.lithosim(mask_pred_bin_orig_size, 0.225, kernels, weight,
#                                                           None, False, avgpool_size=avgpool_size,
#                                                           return_binary_wafer=True, dose=1.02)
#             ilt_loss = ((result - target).pow(4) + (result_inner - target).pow(4) + (result_outer - target).pow(
#                 4)).sum().div(3.0)
#             l2_loss = (bin_mask_outer - bin_mask_inner).abs().sum()  # Image cplx loss
#         else:
#             ilt_loss = (result - target).pow(4).sum()
#             l2_loss = (bin_mask - target).abs().sum()
#
#         new_cord = torch.stack((new_cord[0], new_cord[1], new_cord[2], new_cord[3]), dim=0)
#         cycle_mode = torch.tensor(cycle_mode)
#         cplx_obj = torch.tensor(cplx_obj)
#         # target插值
#         target = torch.nn.functional.interpolate(target, size=(2048, 2048), mode='nearest')
#         ctx.save_for_backward(mask_pred_sig_backup, target, kernels, kernels_ct, kernel_def, kernel_def_ct, weight,
#                               weight_def, new_cord, cycle_mode, cplx_obj, torch.tensor(report_epe))
#
#         if report_epe:
#             checkpoints = get_epe_checkpoints((target.detach().data.cpu().numpy()[0][0] * 255).astype(np.uint8))
#             epe_violation = report_epe_violations((bin_mask.detach().data.cpu().numpy()[0][0] * 255).astype(np.uint8),
#                                                   checkpoints)
#             epe_violation = torch.tensor(epe_violation, requires_grad=False)
#         place_holder = -1
#         place_holder = torch.tensor(place_holder)
#
#         if output_litho_l2_loss:
#             # l2_loss is easier for us to monitor the training and on-nn-ilt correction, it is NOT the exact forward loss of ilt_loss_layer
#             if report_epe:
#                 return l2_loss, epe_violation
#             return l2_loss, place_holder
#         else:
#             # The exact ILT forward loss
#             return ilt_loss, place_holder
#
#     @staticmethod
#     def backward(ctx, grad_output, place_holder):
#         mask_pred_sig, target, kernels, kernels_ct, kernel_def, kernel_def_ct, weight, weight_def, new_cord, _, cplx_obj, _ = ctx.saved_tensors
#         # print(target.shape)
#         new_cord = [new_cord[0], new_cord[1], new_cord[2], new_cord[3]]
#
#         grad_input = ut.compute_gradient_scale(mask_pred_sig, target, kernels, kernels_ct, kernel_def, kernel_def_ct,
#                                                 weight, weight_def, new_cord, cplx_obj.item())
#         grad_input = grad_input * 2  # NOTE: backward for bit_mask_to_two_value_mask function
#
#         return grad_output * grad_input, None, None, None, None, None, None, None, None, None, None, None